package com.gitee.huanminabc.utils_tools.redis.service.impl;

import com.gitee.huanminabc.utils_tools.redis.model.bo.RequestInformationBo;
import com.gitee.huanminabc.utils_tools.redis.redisTemplate.utils.RedisTemplateUtil;
import com.gitee.huanminabc.utils_tools.redis.service.RedisRateLimiterService;
import com.google.common.collect.Lists;
import org.apache.logging.log4j.util.Strings;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Service;

import java.util.ArrayList;

/**
 * @author huanmin
 * @date 2024/2/6
 */
@Service
public class RedisRateLimiterServiceImpl implements RedisRateLimiterService {
    private static final String RATE_LIMITER_KEY_PREFIX="rate_limiter";
    /**
     * 使用redis做限流处理使用的lua脚本
     */
    private static  final String LIMITER_LUA=
            "local key = KEYS[1]\n" +
                    "local count = tonumber(ARGV[1])\n" +
                    "local time = tonumber(ARGV[2])\n" +
                    "local current = redis.call('get', key)\n" +
                    "if current and tonumber(current) > count then\n" +
                    "    return 1\n" +
                    "end\n" +
                    "current = redis.call('incr', key)\n" +
                    "if tonumber(current) == 1 then\n" +
                    "    redis.call('expire', key, time)\n" +
                    "end\n" +
                    "return 0\n";

    /**
     * @param requestInfo 请求信息
     * @return 当前请求是否允许通过
     */
    @Override
    public boolean pass(RequestInformationBo requestInfo) {
        //拿到限流key
        String limiterKey=getRateLimiterKey(requestInfo);
        //执行lua脚本
        RedisScript<Long> redisScript = RedisScript.of(LIMITER_LUA, Long.class);
        Long limiterRes = RedisTemplateUtil.getRedisTemplate().execute(redisScript, Lists.newArrayList(limiterKey), requestInfo.getCount(), requestInfo.getTime());
        //判断限流结果
        return limiterRes==0L;
    }


    private String getRateLimiterKey(RequestInformationBo requestInfo) {
        return combineKey(RATE_LIMITER_KEY_PREFIX,
                requestInfo.getKey(),
                requestInfo.getIp(),
                requestInfo.getMethod().getClass().getName(),
                requestInfo.getMethod().getName());
    }

    private String combineKey(String ... keys) {
        StringBuilder keyBuilder=new StringBuilder();
        for (int i = 0; i < keys.length; i++) {
            if(Strings.isEmpty(keys[i])){
                continue;
            }
            keyBuilder.append(keys[i]);
            if(i==keys.length-1){
                continue;
            }
            keyBuilder.append(":");
        }
        return keyBuilder.toString();
    }

}
